#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EEG实时数据连接器

本文件专门用于与eeg_data_logger.py对接，实现真正的实时数据可视化。
通过继承EEGDataLogger类，重写回调函数来获取实时脑电数据，
并将数据传递给可视化器进行实时图表显示。

使用方法：
1. 确保NeuroSky设备已连接
2. 运行此脚本启动实时监测
3. 实时查看Gamma波等脑电信号变化

原理：
1. 继承EEGDataLogger类，获得完整的数据采集能力
2. 重写所有回调函数，将数据传递给可视化器
3. 使用多线程分离数据采集和图表显示
4. 实现真正的实时数据流可视化
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'eeg'))

import matplotlib
# 设置matplotlib后端，确保图表能够显示
try:
    matplotlib.use('TkAgg')  # 使用TkAgg后端
except ImportError:
    try:
        matplotlib.use('Qt5Agg')  # 备选Qt5Agg后端
    except ImportError:
        matplotlib.use('Agg')  # 最后使用Agg后端
        print("警告: 使用Agg后端，图表可能无法显示")

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.dates import DateFormatter
import numpy as np
from datetime import datetime, timedelta
from collections import deque
import threading
import time

try:
    from eeg_data_logger import EEGDataLogger
except ImportError:
    print("错误: 无法导入eeg_data_logger模块")
    print("请确保eeg_data_logger.py文件在正确的路径下")
    sys.exit(1)

class EEGLiveConnector(EEGDataLogger):
    """EEG实时数据连接器类"""
    
    def __init__(self, port="COM5", baudrate=57600, window_size=30):
        """
        初始化EEG实时数据连接器
        
        Args:
            port (str): 串口号，默认COM5
            baudrate (int): 波特率，默认57600
            window_size (int): 显示窗口大小（秒），默认30秒
        """
        # 初始化父类，不需要日志文件
        super().__init__(port, baudrate, log_file=None)
        
        self.window_size = window_size
        
        # 数据缓存队列
        self.max_points = window_size * 10  # 假设每秒10个数据点
        self.time_data = deque(maxlen=self.max_points)
        self.attention_data = deque(maxlen=self.max_points)
        self.meditation_data = deque(maxlen=self.max_points)
        self.delta_data = deque(maxlen=self.max_points)
        self.theta_data = deque(maxlen=self.max_points)
        self.low_alpha_data = deque(maxlen=self.max_points)
        self.high_alpha_data = deque(maxlen=self.max_points)
        self.low_beta_data = deque(maxlen=self.max_points)
        self.high_beta_data = deque(maxlen=self.max_points)
        self.low_gamma_data = deque(maxlen=self.max_points)
        self.mid_gamma_data = deque(maxlen=self.max_points)
        self.signal_quality_data = deque(maxlen=self.max_points)
        
        # 图表相关
        self.fig = None
        self.axes = None
        self.lines = {}
        self.ani = None
        
        # 线程控制
        self.data_lock = threading.Lock()
        self.is_visualizing = False
        
        # 统计信息
        self.data_count = 0
        self.start_time = None
    
    def _init_csv_file(self):
        """重写父类方法，不需要CSV文件"""
        pass
    
    def _log_data_to_csv(self):
        """重写父类方法，不需要CSV记录"""
        pass
    
    def _add_data_point(self):
        """添加数据点到缓存队列"""
        timestamp = datetime.now()
        
        with self.data_lock:
            self.time_data.append(timestamp)
            self.attention_data.append(self.current_data['attention'])
            self.meditation_data.append(self.current_data['meditation'])
            self.delta_data.append(self.current_data['delta'])
            self.theta_data.append(self.current_data['theta'])
            self.low_alpha_data.append(self.current_data['lowAlpha'])
            self.high_alpha_data.append(self.current_data['highAlpha'])
            self.low_beta_data.append(self.current_data['lowBeta'])
            self.high_beta_data.append(self.current_data['highBeta'])
            self.low_gamma_data.append(self.current_data['lowGamma'])
            self.mid_gamma_data.append(self.current_data['midGamma'])
            # 转换信号质量：poorSignal越小，质量越好
            signal_quality = 200 - self.current_data['poorSignal']
            self.signal_quality_data.append(signal_quality)
            
            self.data_count += 1
            if self.start_time is None:
                self.start_time = timestamp
    
    # 重写所有回调函数，添加数据到可视化缓存
    def attention_callback(self, value):
        """注意力数据回调函数"""
        self.current_data['attention'] = value
        print(f"注意力: {value}")
        self._add_data_point()
    
    def meditation_callback(self, value):
        """冥想数据回调函数"""
        self.current_data['meditation'] = value
        print(f"冥想: {value}")
        self._add_data_point()
    
    def delta_callback(self, value):
        """Delta波回调函数"""
        self.current_data['delta'] = value
        print(f"Delta波: {value}")
        self._add_data_point()
    
    def theta_callback(self, value):
        """Theta波回调函数"""
        self.current_data['theta'] = value
        print(f"Theta波: {value}")
        self._add_data_point()
    
    def low_alpha_callback(self, value):
        """低Alpha波回调函数"""
        self.current_data['lowAlpha'] = value
        print(f"低Alpha波: {value}")
        self._add_data_point()
    
    def high_alpha_callback(self, value):
        """高Alpha波回调函数"""
        self.current_data['highAlpha'] = value
        print(f"高Alpha波: {value}")
        self._add_data_point()
    
    def low_beta_callback(self, value):
        """低Beta波回调函数"""
        self.current_data['lowBeta'] = value
        print(f"低Beta波: {value}")
        self._add_data_point()
    
    def high_beta_callback(self, value):
        """高Beta波回调函数"""
        self.current_data['highBeta'] = value
        print(f"高Beta波: {value}")
        self._add_data_point()
    
    def low_gamma_callback(self, value):
        """低Gamma波回调函数"""
        self.current_data['lowGamma'] = value
        print(f"低Gamma波: {value}")
        self._add_data_point()
    
    def mid_gamma_callback(self, value):
        """中Gamma波回调函数"""
        self.current_data['midGamma'] = value
        print(f"中Gamma波: {value}")
        self._add_data_point()
    
    def poor_signal_callback(self, value):
        """信号质量回调函数"""
        self.current_data['poorSignal'] = value
        print(f"信号质量: {value}")
        self._add_data_point()
    
    def setup_plots(self):
        """设置图表布局"""
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 创建图表 - 3x2布局，专注于Gamma波，增大图表尺寸
        self.fig, self.axes = plt.subplots(3, 2, figsize=(20, 15))
        self.fig.suptitle('实时脑电信号监测 - 与eeg_data_logger.py对接', fontsize=16, fontweight='bold', color='darkblue')
        
        # 子图1: 注意力和冥想状态
        ax1 = self.axes[0, 0]
        ax1.set_title('注意力 & 冥想状态', fontweight='bold')
        ax1.set_ylabel('指数值 (0-100)')
        ax1.set_ylim(0, 100)
        ax1.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        self.lines['attention'], = ax1.plot([], [], 'r-', label='注意力', linewidth=3)
        self.lines['meditation'], = ax1.plot([], [], 'b-', label='冥想', linewidth=3)
        ax1.legend(loc='upper right', framealpha=0.9)
        ax1.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图2: Gamma波（重点显示）
        ax2 = self.axes[0, 1]
        ax2.set_title('Gamma波实时监测', fontweight='bold', color='red')
        ax2.set_ylabel('功率值')
        ax2.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        # self.lines['low_gamma'], = ax2.plot([], [], 'g-', label='低Gamma', linewidth=3)
        # self.lines['mid_gamma'], = ax2.plot([], [], 'm-', label='中Gamma', linewidth=3)
        # 修改Gamma波的绘制方式，明确指定为折线图
        # 在setup_plots函数中，为所有线条明确指定折线样式
        self.lines['total_gamma'], = ax2.plot([], [], color='#FF4500', linestyle='-', linewidth=3, label='总Gamma', alpha=0.9, solid_capstyle='round', solid_joinstyle='round')
        ax2.legend(loc='upper right', framealpha=0.9)
        ax2.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图3: Alpha波段
        ax3 = self.axes[1, 0]
        ax3.set_title('Alpha波段', fontweight='bold')
        ax3.set_ylabel('功率值')
        ax3.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        self.lines['low_alpha'], = ax3.plot([], [], '#FF8C00', label='低Alpha', linewidth=3)
        self.lines['high_alpha'], = ax3.plot([], [], '#FFD700', label='高Alpha', linewidth=3)
        ax3.legend(loc='upper right', framealpha=0.9)
        ax3.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图4: Beta波段
        ax4 = self.axes[1, 1]
        ax4.set_title('Beta波段', fontweight='bold')
        ax4.set_ylabel('功率值')
        ax4.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        self.lines['low_beta'], = ax4.plot([], [], '#00CED1', label='低Beta', linewidth=3)
        self.lines['high_beta'], = ax4.plot([], [], '#000080', label='高Beta', linewidth=3)
        ax4.legend(loc='upper right', framealpha=0.9)
        ax4.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图5: Delta & Theta波段
        ax5 = self.axes[2, 0]
        ax5.set_title('Delta & Theta波段', fontweight='bold')
        ax5.set_ylabel('功率值')
        ax5.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        self.lines['delta'], = ax5.plot([], [], '#8A2BE2', label='Delta', linewidth=3)
        self.lines['theta'], = ax5.plot([], [], '#A0522D', label='Theta', linewidth=3)
        ax5.legend(loc='upper right', framealpha=0.9)
        ax5.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图6: 信号质量和实时统计
        ax6 = self.axes[2, 1]
        ax6.set_title('信号质量 & 实时统计', fontweight='bold')
        ax6.set_ylabel('信号质量 (0-200)')
        ax6.set_ylim(0, 200)
        ax6.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        self.lines['signal_quality'], = ax6.plot([], [], '#2F4F4F', label='信号质量', linewidth=3)
        ax6.legend(loc='upper right', framealpha=0.9)
        ax6.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 添加实时统计显示文本
        self.text_display = ax6.text(0.02, 0.98, '', transform=ax6.transAxes, 
                                   fontsize=8, verticalalignment='top',
                                   bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8))
        
        plt.tight_layout()
        plt.subplots_adjust(top=0.95)
    
    def update_plots(self, frame):
        """更新图表数据"""
        with self.data_lock:
            if len(self.time_data) == 0:
                return list(self.lines.values())
            
            # 获取当前数据
            times = list(self.time_data)
            attention = list(self.attention_data)
            meditation = list(self.meditation_data)
            delta = list(self.delta_data)
            theta = list(self.theta_data)
            low_alpha = list(self.low_alpha_data)
            high_alpha = list(self.high_alpha_data)
            low_beta = list(self.low_beta_data)
            high_beta = list(self.high_beta_data)
            low_gamma = list(self.low_gamma_data)
            mid_gamma = list(self.mid_gamma_data)
            signal_quality = list(self.signal_quality_data)
        
        # 计算总Gamma
        total_gamma = [lg + mg for lg, mg in zip(low_gamma, mid_gamma)]
        
        # 更新所有线条数据
        self.lines['attention'].set_data(times, attention)
        self.lines['meditation'].set_data(times, meditation)
        # self.lines['low_gamma'].set_data(times, low_gamma)
        # self.lines['mid_gamma'].set_data(times, mid_gamma)
        self.lines['total_gamma'].set_data(times, total_gamma)
        self.lines['low_alpha'].set_data(times, low_alpha)
        self.lines['high_alpha'].set_data(times, high_alpha)
        self.lines['low_beta'].set_data(times, low_beta)
        self.lines['high_beta'].set_data(times, high_beta)
        self.lines['delta'].set_data(times, delta)
        self.lines['theta'].set_data(times, theta)
        self.lines['signal_quality'].set_data(times, signal_quality)
        
        # 动态调整x轴范围
        if len(times) > 0:
            current_time = times[-1]
            start_time = current_time - timedelta(seconds=self.window_size)
            
            for ax in self.axes.flat:
                ax.set_xlim(start_time, current_time)
        
        # 动态调整Y轴范围，便于观察变化趋势
        # 注意力和冥想状态已经固定在0-100
        # 其他频段使用动态范围，根据实际数据调整
        if len(times) > 0:
            # Gamma波子图 - 使用动态范围
            if total_gamma:  # 只检查总Gamma数据
                gamma_max = max(total_gamma) if total_gamma else 0
                if gamma_max > 0:
                    self.axes[0, 1].set_ylim(0, gamma_max * 1.2)  # 留20%余量
                else:
                    self.axes[0, 1].set_ylim(0, 1000)  # 默认小范围
            
            # Alpha波子图 - 使用动态范围
            if low_alpha or high_alpha:
                alpha_max = max(max(low_alpha) if low_alpha else 0,
                              max(high_alpha) if high_alpha else 0)
                if alpha_max > 0:
                    self.axes[1, 0].set_ylim(0, alpha_max * 1.2)
                else:
                    self.axes[1, 0].set_ylim(0, 10000)
            
            # Beta波子图 - 使用动态范围
            if low_beta or high_beta:
                beta_max = max(max(low_beta) if low_beta else 0,
                             max(high_beta) if high_beta else 0)
                if beta_max > 0:
                    self.axes[1, 1].set_ylim(0, beta_max * 1.2)
                else:
                    self.axes[1, 1].set_ylim(0, 5000)
            
            # Delta & Theta子图 - 使用动态范围
            if delta or theta:
                dt_max = max(max(delta) if delta else 0,
                           max(theta) if theta else 0)
                if dt_max > 0:
                    self.axes[2, 0].set_ylim(0, dt_max * 1.2)
                else:
                    self.axes[2, 0].set_ylim(0, 50000)
        
        # 更新实时统计显示
        if len(times) > 0:
            current_values = {
                'attention': attention[-1] if attention else 0,
                'meditation': meditation[-1] if meditation else 0,
                'low_gamma': low_gamma[-1] if low_gamma else 0,
                'mid_gamma': mid_gamma[-1] if mid_gamma else 0,
                'signal_quality': signal_quality[-1] if signal_quality else 0
            }
            
            # 计算运行时间
            if self.start_time:
                runtime = times[-1] - self.start_time
                runtime_str = str(runtime).split('.')[0]  # 去掉微秒
            else:
                runtime_str = "00:00:00"
            
            # 计算数据速率
            data_rate = self.data_count / max(1, (times[-1] - self.start_time).total_seconds()) if self.start_time else 0
            
            text_content = f"""实时统计:
注意力: {current_values['attention']:.0f}
冥想: {current_values['meditation']:.0f}
低Gamma: {current_values['low_gamma']:.0f}
中Gamma: {current_values['mid_gamma']:.0f}
总Gamma: {current_values['low_gamma'] + current_values['mid_gamma']:.0f}
信号质量: {current_values['signal_quality']:.0f}

运行时间: {runtime_str}
数据点数: {self.data_count}
数据速率: {data_rate:.1f}/秒"""
            
            self.text_display.set_text(text_content)
        
        return list(self.lines.values())
    
    def start_visualization(self):
        """开始实时可视化"""
        try:
            # 设置图表
            self.setup_plots()
            
            # 启动数据采集线程
            self.data_thread = threading.Thread(target=self._start_data_collection)
            self.data_thread.daemon = True
            self.data_thread.start()
            
            # 启动动画
            self.ani = animation.FuncAnimation(
                self.fig, self.update_plots, interval=100,
                blit=False, cache_frame_data=False
            )
            
            print("=== EEG实时数据可视化已启动 ===")
            print("与eeg_data_logger.py完全对接")
            print("关闭图表窗口或按Ctrl+C停止可视化")
            
            # 显示图表
            print("正在启动图表显示...")
            print("如果图表窗口没有出现，请检查matplotlib后端设置")
            plt.show(block=True)  # 阻塞显示，确保窗口不会立即关闭
            
        except KeyboardInterrupt:
            print("\n用户中断，停止可视化...")
            self.stop_visualization()
        except Exception as e:
            print(f"可视化错误: {e}")
            self.stop_visualization()
    
    def _start_data_collection(self):
        """启动数据采集（在单独线程中运行）"""
        try:
            # 初始化NeuroPy对象
            from NeuroPy3 import NeuroPy
            self.neuropy = NeuroPy(self.port, self.baudrate)
            
            # 设置所有回调函数
            self.neuropy.setCallBack("attention", self.attention_callback)
            self.neuropy.setCallBack("meditation", self.meditation_callback)
            self.neuropy.setCallBack("delta", self.delta_callback)
            self.neuropy.setCallBack("theta", self.theta_callback)
            self.neuropy.setCallBack("lowAlpha", self.low_alpha_callback)
            self.neuropy.setCallBack("highAlpha", self.high_alpha_callback)
            self.neuropy.setCallBack("lowBeta", self.low_beta_callback)
            self.neuropy.setCallBack("highBeta", self.high_beta_callback)
            self.neuropy.setCallBack("lowGamma", self.low_gamma_callback)
            self.neuropy.setCallBack("midGamma", self.mid_gamma_callback)
            self.neuropy.setCallBack("poorSignal", self.poor_signal_callback)
            
            # 开始数据采集
            self.is_visualizing = True
            self.neuropy.start()
            
            print(f"数据采集已启动 - 串口: {self.port}, 波特率: {self.baudrate}")
            
            # 保持数据采集线程运行
            while self.is_visualizing:
                time.sleep(0.1)
                
        except Exception as e:
            print(f"数据采集错误: {e}")
            print("请检查：")
            print("1. NeuroSky设备是否正确连接")
            print("2. 串口号是否正确")
            print("3. 设备驱动是否已安装")
    
    def stop_visualization(self):
        """停止实时可视化"""
        self.is_visualizing = False
        if self.neuropy:
            self.neuropy.stop()
        if self.ani:
            self.ani.event_source.stop()
        plt.close('all')
        print("EEG实时数据可视化已停止")

def main():
    """主函数"""
    print("=== EEG实时数据连接器 ===")
    print("与eeg_data_logger.py完全对接的实时可视化工具")
    print("专门用于实时监测Gamma波等脑电信号")
    print()
    
    # 询问用户串口设置
    port = input("请输入串口号 (默认COM5): ").strip() or "COM5"
    baudrate_input = input("请输入波特率 (默认57600): ").strip()
    baudrate = int(baudrate_input) if baudrate_input else 57600
    
    print(f"使用串口: {port}, 波特率: {baudrate}")
    print("请确保NeuroSky设备已正确连接...")
    
    # 创建连接器实例
    connector = EEGLiveConnector(
        port=port,
        baudrate=baudrate,
        window_size=30  # 显示最近30秒的数据
    )
    
    # 启动实时可视化
    connector.start_visualization()

if __name__ == "__main__":
    main()